package com.test.html;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import okhttp3.*;
import okio.BufferedSink;
import okio.Okio;
import okio.Sink;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 下载Lexica.art图片
 */
public class Lexica {
    private static AtomicInteger handled = new AtomicInteger(0);

    private static LocalDate now = LocalDate.now();
    private static String baseDir = "D:\\tmp\\lexica\\" + now.format(DateTimeFormatter.ofPattern("yy-MM"));
    private static String extFix = ".jpg";

    public static void main(String[] args) throws IOException {
        System.out.println(baseDir);
        Path path = Paths.get(baseDir);
        Files.createDirectories(path);

        /* 下载单个图片
        download("https://cdn.midjourney.com/1b9f6d82-a5dd-4600-98d8-3a54c637a7bd/0_1.png");
        try {
            Thread.sleep(4000L);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }*/

        // 从midjourney下载图片
        String url = "https://lexica.art/api/infinite-prompts";
        String imageUrl = "https://image.lexica.art/md2/";
        //  查询图片的描述
        String text = "superwomen";
        // 查询图片的起始位置，需累加
        Integer cursor = 200;

//        String doc = OkHttpUtilsLocal.builder()
//                .url(url)
//                .post(true)
//                .sync();
        OkHttpClient client = new OkHttpClient().newBuilder()
                .build();
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, "{\"text\":\"" + text + "\",\"searchMode\":\"images\",\"source\":\"search\",\"cursor\":" + cursor + ",\"model\":\"lexica-aperture-v2\"}");
        Request request = new Request.Builder()
                .url(url)
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                // .addHeader("Cookie", "__Host-next-auth.csrf-token=a812c486e16e13736a67070ebf8e52fb1db6b8eb21f16a0cd7b2fe257e1d17a4%7Ceeff6ca92c9aebf5ff0c26fde048085e452368caa591570ac5dd7717f1454a7d; __Secure-next-auth.callback-url=https%3A%2F%2Flexica.art")
                .build();
        Response response = client.newCall(request).execute();
        String  doc = response.body().string();
        // System.out.println(doc);
        JSONObject data = JSONObject.parseObject(doc);
        JSONArray jobs = data.getJSONArray("images");
        int size = jobs.size();
        JSONObject job;
        for (int i = 0; i < size; i++) {
            job = jobs.getJSONObject(i);
            File dest = new File(baseDir, getFileNameFromUrl(imageUrl + job.getString("id")));
            if (dest.exists()) {
                handled.incrementAndGet();
                System.out.println(imageUrl + job.getString("id") + "已存在");
                continue;
            }
            downloadFile(imageUrl + job.getString("id"));
        }
        System.out.println("已有文件数：" + handled.get());
        while (handled.get() < size) {
            try {
                Thread.sleep(1000L);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        System.out.println(handled.get());
        System.exit(1);
    }

    private static String getFileNameFromUrl(String url) {
        return now.format(DateTimeFormatter.ofPattern("yyMM")) + "_" + url.substring(url.lastIndexOf("/") + 1) + extFix;
    }

    private static final OkHttpClient client = new OkHttpClient();
    private static final int REQUEST_EXTERNAL_STORAGE = 1;

    //下载文件
    private static void downloadFile(final String url) {
        new Thread(() -> download(url)).start();
    }

    //下载
    public static void download(final String url) {

        final long startTime = System.currentTimeMillis();
        // System.out.println("DOWNLOAD" + "startTime=" + startTime);
        Request request = new Request.Builder().url(url).build();
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                handled.incrementAndGet();
                // 下载失败
                e.printStackTrace();
                System.out.println("DOWNLOAD" + "download failed");
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                Sink sink = null;
                BufferedSink bufferedSink = null;
                try {
                    File dest = new File(baseDir, getFileNameFromUrl(url)); // url.substring(27, url.lastIndexOf("/")) + ".png"
                    if (dest.exists()) {
                        System.out.println(url + "已存在");
                        return;
                    }
                    sink = Okio.sink(dest);
                    bufferedSink = Okio.buffer(sink);
                    bufferedSink.writeAll(response.body().source());

                    bufferedSink.close();
                    System.out.println(url + " download success " + "totalTime=" + (System.currentTimeMillis() - startTime));
                } catch (Exception e) {
                    e.printStackTrace();
                    System.out.println("DOWNLOAD" + "download failed");
                } finally {
                    handled.incrementAndGet();
                    if (bufferedSink != null) {
                        bufferedSink.close();
                    }
                    response.close();
                }
            }
        });
    }
}
